# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2.13
import csv
import re
import copy
import pandas as pd
import numpy as np
from typing import List, Tuple

from utils.instances import TOKENIZER, LLM
from utils import prompts
from langchain_core.prompts import ChatPromptTemplate
import utils.configFinRAG as configFinRAG


def calculate_jaccard_similarity(tokens: List[int], token_lists: List[List[int]]) -> List[float]:
    """
    计算Jaccard相似度
    :param tokens: 目标问题的token列表
    :param token_lists: 所有示例问题的token列表
    :return: 相似度列表
    """
    similarity_list = []
    for example_tokens in token_lists:
        intersection = len(set(tokens) & set(example_tokens))
        union = len(set(tokens)) + len(set(example_tokens))
        similarity_list.append(intersection / union if union > 0 else 0)
    return similarity_list


def get_top_n_indices(similarity_list: List[float], n: int) -> List[int]:
    """
    获取相似度最高的n个索引
    :param similarity_list: 相似度列表
    :param n: 需要返回的索引数量
    :return: 相似度最高的n个索引
    """
    return np.argsort(similarity_list)[-n:][::-1].tolist()


def extract_sql_from_response(response: str) -> str:
    """
    从模型响应中提取SQL语句
    :param response: 模型生成的响应
    :return: 提取的SQL语句
    """
    sql_pattern = r'```sql(.*?)```'
    match = re.search(sql_pattern, response, re.DOTALL)
    return match.group(1).strip() if match else "error"


def generate_sql(question: str, llm, example_question_list: List[str], example_sql_list: List[str],
                tmp_example_token_list: List[List[int]], example_num: int = 5) -> Tuple[str, str]:
    """
    生成SQL语句
    :param question: 用户问题
    :param llm: 大语言模型
    :param example_question_list: 示例问题列表
    :param example_sql_list: 示例SQL列表
    :param tmp_example_token_list: 示例问题的token列表
    :param example_num: 示例数量
    :return: (prompt, SQL语句)
    """
    # 提取日期并替换为空格
    date_list = re.findall(r'\d{8}', question)
    temp_question_for_search = re.sub(r'\d{8}', ' ', question)

    # 计算Jaccard相似度
    temp_tokens = TOKENIZER(temp_question_for_search)['input_ids']
    similarity_list = calculate_jaccard_similarity(temp_tokens, tmp_example_token_list)

    # 获取相似度最高的n个索引
    top_indices = get_top_n_indices(similarity_list, example_num)

    # 组装prompt
    examples = ""
    for index in top_indices:
        examples += f"问题：{example_question_list[index]}\nSQL：{example_sql_list[index]}\n"
        if len(examples) > 2000:  # 防止提示语过长
            break

    # 调用模型生成SQL
    prompt = ChatPromptTemplate.from_template(prompts.GENERATE_SQL_TEMPLATE)
    chain = prompt | llm
    response = chain.invoke({"examples": examples, "table_info": prompts.TABLE_INFO, "question": question})

    # 提取SQL语句
    sql = extract_sql_from_response(response.content)
    if sql == "error":
        print("generate sql error:", question)
        return "error", "error"
    return prompt.invoke({"examples": examples, "table_info": prompts.TABLE_INFO, "question": question}), sql


def main():
    # 读取问题和SQL模板
    sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
    example_question_list = sql_examples_file['问题'].tolist()
    example_sql_list = sql_examples_file['SQL'].tolist()
    example_token_list = [TOKENIZER(question)['input_ids'] for question in example_question_list]

    # 读取测试问题
    question_csv_file = pd.read_csv(configFinRAG.question_classify_path, delimiter=",", header=0)

    # 打开结果文件
    with open(configFinRAG.question_sql_path, 'w', newline='', encoding='utf-8-sig') as question_sql_file:
        csvwriter = csv.writer(question_sql_file)
        csvwriter.writerow(['问题id', '问题', 'SQL', 'prompt'])

        # 循环处理问题
        for _, row in question_csv_file.iterrows():
            if row['分类'] == '查询数据库':
                result_prompt, result = generate_sql(row['问题'], LLM, example_question_list, example_sql_list,
                                                     example_token_list)
                csvwriter.writerow([row['问题id'], row['问题'], result, result_prompt])
            else:
                print("pass question:", row['问题'])


if __name__ == '__main__':
    main()